package org.example.utils;

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.opencv.OpenCVImageFactory;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import cn.hutool.http.HttpStatus;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.litongjava.ai.djl.paddle.ocr.v4.common.RotatedBox;
import com.litongjava.ai.djl.paddle.ocr.v4.common.RotatedBoxCompX;
import com.litongjava.ai.djl.paddle.ocr.v4.detection.OcrV4Detection;
import com.litongjava.ai.djl.paddle.ocr.v4.recognition.OcrV4Recognition;
import net.sourceforge.tess4j.ITessAPI;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.Word;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang.StringUtils;
import org.example.model.PaddleWidget;
import org.example.model.TrWidget;
import org.example.model.Tess4JWidget;
import org.example.model.Widget;
import org.openqa.selenium.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.Rectangle;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.math.BigDecimal;
import java.util.*;
import java.util.List;

public class OcrUtils {

    private static Logger logger = LoggerFactory.getLogger(OcrUtils.class);

    private static ITesseract iTesseract;

    private static ImageFactory imageFactory;

    private static OcrV4Detection detection;

    private static OcrV4Recognition recognition;

    private static String TR_URL;

    private static Map<String, Object> params;

    private static ZooModel detectionModel;
    private static Predictor<ai.djl.modality.cv.Image, NDList> detector;
    private static ZooModel recognitionModel;
    private static Predictor<ai.djl.modality.cv.Image, String> recognizer;
    private static NDManager manager;

    public static final String OCR_TYPE_TESS4J = "tess4j";

    public static final String OCR_TYPE_TR = "tr";

    public static final String OCR_TYPE_PADDLE = "paddle";

    private static String OCR_TYPE = "";
    /**
     * 初始化OCR
     * @param type tess4j 、 tr 、 paddle
     * @param params
     *      tess4j 需要传入两个参数：数据路径、需要识别的语言类型
     *      tr 需要传入一个参数：tr的url路径
     *      paddle 不需要传入参数
     */
    public static void initOcr(String type, String ... params) {
        OCR_TYPE = type;
        if (OCR_TYPE_TESS4J.equals(type)) {
            initOcrTess4J(params[0], params[1]);
        } else if (OCR_TYPE_TR.equals(type)) {
            initOcrTr(params[0]);
        } else if (OCR_TYPE_PADDLE.equals(type)) {
            initOcrPaddle();
        } else {
            logger.warn("ocr 初始化参数传入错误！");
        }
    }

    /**
     * 识别文字
     * @param webDriver
     * @param saveFilePath
     * @return
     */
    public static List<Widget> ocrImg(WebDriver webDriver, String saveFilePath) {
        List<Widget> widgetList = new ArrayList<>();
        if (OCR_TYPE_TESS4J.equals(OCR_TYPE)) {
            widgetList = ocrImgTess4J(webDriver, saveFilePath);
        } else if (OCR_TYPE_TR.equals(OCR_TYPE)) {
            widgetList = ocrImgTr(webDriver, saveFilePath);
        } else if (OCR_TYPE_PADDLE.equals(OCR_TYPE)) {
            widgetList = ocrImgPaddle(webDriver, saveFilePath);
        } else {
            logger.warn("ocr 初始化参数传入错误！");
        }

        return widgetList;
    }

    /**
     * 识别文字
     * @param file
     * @return
     */
    public static List<Widget> ocrImg(File file) {
        List<Widget> widgetList = new ArrayList<>();
        if (OCR_TYPE_TESS4J.equals(OCR_TYPE)) {
            widgetList = ocrImgTess4J(file);
        } else if (OCR_TYPE_TR.equals(OCR_TYPE)) {
            widgetList = ocrImgTr(file);
        } else if (OCR_TYPE_PADDLE.equals(OCR_TYPE)) {
            widgetList = ocrImgPaddle(file);
        } else {
            logger.warn("ocr 初始化参数传入错误！");
        }

        return widgetList;
    }

    /**
     * 释放ocr资源
     */
    public static void release() {
        if (OCR_TYPE_TESS4J.equals(OCR_TYPE)) {
            // 无需释放
        } else if (OCR_TYPE_TR.equals(OCR_TYPE)) {
            // 无需释放
        } else if (OCR_TYPE_PADDLE.equals(OCR_TYPE)) {
            releasePaddle();
        } else {
            logger.warn("ocr 初始化参数传入错误！");
        }
    }
    /**
     * 初始化ocr
     * @param language 需要识别的语种
     */
    private static void initOcrTess4J(String dataPath, String language) {
        iTesseract = new Tesseract();
        // 设置data目录
        iTesseract.setDatapath(dataPath);
        // 设置语言
        iTesseract.setLanguage(language);
        // iTesseract.setVariable("user_defined_dpi", "300");
        List<String> configList = new ArrayList<String>();
        // configList.add("digits");
        iTesseract.setConfigs(configList);
        // 设置OCR引擎模式（OEM）
        iTesseract.setOcrEngineMode(1);
        // 设置图片分割模式（PSM）
        iTesseract.setPageSegMode(6);
    }

    private static void initOcrPaddle() {
        imageFactory = OpenCVImageFactory.getInstance();
        detection = new OcrV4Detection();
        recognition = new OcrV4Recognition();
        try {
            detectionModel = ModelZoo.loadModel(detection.chDetCriteria());
            detector = detectionModel.newPredictor();
            recognitionModel = ModelZoo.loadModel(recognition.chRecCriteria());
            recognizer = recognitionModel.newPredictor();
            manager = NDManager.newBaseManager();
        } catch (Exception e) {
            logger.error("paddle ocr init error:" + e);
        }
    }

    private static void initOcrTr(String trUrl) {
        // TR_URL = "http://82.156.132.121:8089/api/tr-run/";
        TR_URL = trUrl;
        params = new HashMap<>();
        params.put("is_draw", 0);
        params.put("compress", 0);
    }


    private static List<Widget> ocrImgTess4J(WebDriver webDriver, String saveFilePath) {
        List<Widget> list = new ArrayList<>();
        try {
            WebElement canvas = webDriver.findElement(By.xpath("//*[@id=\"content\"]/div[2]/div[2]/div/div/div[2]/div[1]"));
            File screenshotFile = canvas.getScreenshotAs(OutputType.FILE);
            if (StringUtils.isNotEmpty(saveFilePath)) {
                try {
                    FileUtils.copyFile(screenshotFile, new File(saveFilePath + GlobalUtils.randomStr() + ".png"));
                } catch (IOException e) {
                    GlobalUtils.logi("存储路径有误！" + e);
                }
            }
            Image image = ImageIO.read(screenshotFile);
            BufferedImage bi = toBufferedImage(image);

            for(Word word : iTesseract.getWords(bi, ITessAPI.TessPageIteratorLevel.RIL_TEXTLINE)) {
                Rectangle rect = word.getBoundingBox();
                Tess4JWidget tess4JWidget = new Tess4JWidget();
                tess4JWidget.setText(word.getText());
                tess4JWidget.setX1(rect.getMinX());
                tess4JWidget.setX2(rect.getMaxX());
                tess4JWidget.setY1(rect.getMinY());
                tess4JWidget.setY2(rect.getMaxY());
                list.add(tess4JWidget);
            }
        } catch (Exception e) {
            logger.error("ocr识图出错：" + e);
        }
        return list;
    }

    private static List<Widget> ocrImgTess4J(File file) {
        List<Widget> list = new ArrayList<>();
        try {
            Image image = ImageIO.read(file);
            BufferedImage bi = toBufferedImage(image);

            for(Word word : iTesseract.getWords(bi, ITessAPI.TessPageIteratorLevel.RIL_TEXTLINE)) {
                Rectangle rect = word.getBoundingBox();
                Tess4JWidget tess4JWidget = new Tess4JWidget();
                tess4JWidget.setText(word.getText());
                tess4JWidget.setX1(rect.getMinX());
                tess4JWidget.setX2(rect.getMaxX());
                tess4JWidget.setY1(rect.getMinY());
                tess4JWidget.setY2(rect.getMaxY());
                list.add(tess4JWidget);
            }
        } catch (Exception e) {
            logger.error("ocr识图出错：" + e);
        }
        return list;
    }

    /**
     * Tr识别返回文字
     * @param webDriver
     * @param saveFilePath
     * @return
     */
    public static List<Widget> ocrImgTr(WebDriver webDriver, String saveFilePath) {
        WebElement canvas = webDriver.findElement(By.xpath("//*[@id=\"content\"]/div[2]/div[2]/div/div/div[2]/div[1]"));
        File screenshotFile = canvas.getScreenshotAs(OutputType.FILE);
        if (StringUtils.isNotEmpty(saveFilePath)) {
            try {
                FileUtils.copyFile(screenshotFile, new File(saveFilePath + GlobalUtils.randomStr() + ".png"));
            } catch (IOException e) {
                GlobalUtils.logi("存储路径有误！" + e);
            }
        }
        try {
            return convertTrTextRes(HttpUtils.requestOCRForHttp(TR_URL, params, screenshotFile));
        } catch (Exception e) {
            logger.error("ocr识图出错：" + e);
        }
        return null;
    }

    /**
     * Tr识别返回文字
     * @param file
     * @return
     */
    public static List<Widget> ocrImgTr(File file) {
        try {
            return convertTrTextRes(HttpUtils.requestOCRForHttp(TR_URL, params, file));
        } catch (Exception e) {
            logger.error("ocr识图出错：" + e);
        }
        return null;
    }

    /**
     * paddle识别文字
     * @param webDriver
     * @param saveFilePath
     * @return
     */
    public static List<Widget> ocrImgPaddle(WebDriver webDriver, String saveFilePath) {
        WebElement canvas = webDriver.findElement(By.xpath("//*[@id=\"content\"]/div[2]/div[2]/div/div/div[2]/div[1]"));
        File screenshotFile = canvas.getScreenshotAs(OutputType.FILE);
        if (StringUtils.isNotEmpty(saveFilePath)) {
            try {
                FileUtils.copyFile(screenshotFile, new File(saveFilePath + GlobalUtils.randomStr() + ".png"));
            } catch (IOException e) {
                GlobalUtils.logi("存储路径有误！" + e);
            }
        }
        return ocrImgPaddle(screenshotFile);
    }

    /**
     * paddle识别文字
     * @param file
     * @return
     */
    public static List<Widget> ocrImgPaddle(File file) {
        ai.djl.modality.cv.Image image;
        try {
            image = imageFactory.fromInputStream(new FileInputStream(file));
        } catch (IOException e) {
            logger.error("读取文件失败! " + e);
            return null;
        }

        List<Widget> widgetList = new ArrayList<>();
        try {
            long timeInferStart = System.currentTimeMillis();
            List<RotatedBox> detections = recognition.predict(manager, image, detector, recognizer);
            for (int i = 0; i < detections.size(); i++) {
                RotatedBox rotatedBox = detections.get(i);
                NDArray box = rotatedBox.getBox();
                // 8个数值
                float[] pointsArr = box.toFloatArray();

                // 左上
                float[] lt = java.util.Arrays.copyOfRange(pointsArr, 0, 2);

                // 左下
                float[] lb = java.util.Arrays.copyOfRange(pointsArr, 6, 8);

                // 右上
                float[] rt = java.util.Arrays.copyOfRange(pointsArr, 2, 4);

                // 右下
                float[] rb = java.util.Arrays.copyOfRange(pointsArr, 4, 6);
                float width = distance(lt, rt);
                float height = distance(lt, lb);

                // 文字内容
                String text = rotatedBox.getText();

                PaddleWidget widget = new PaddleWidget();
                widget.setText(text);
                widget.setX1(lt[0]);
                widget.setY1(lt[1]);
                widget.setX2(rb[0]);
                widget.setY2(rb[1]);
                widget.setWidth(width);
                widget.setHeight(height);
                widgetList.add(widget);
            }

//            for (int i = 0; i < 1000; i++) {
//                detections = recognition.predict(image, detector, recognizer);
//                for (RotatedBox result : detections) {
//                    System.out.println(result.getText());
//                }
//                System.out.println("index : " + i);
//            }

            long timeInferEnd = System.currentTimeMillis();
            // System.out.println("time: " + (timeInferEnd - timeInferStart));

            // 对检测结果根据坐标位置，根据从上到下，从做到右，重新排序，下面算法对图片倾斜旋转角度较小的情形适用
            // 如果图片旋转角度较大，则需要自行改进算法，需要根据斜率校正计算位置。
            // Reorder the detection results based on the coordinate positions, from top to bottom, from left to right. The algorithm below is suitable for situations where the image is slightly tilted or rotated.
            // If the image rotation angle is large, the algorithm needs to be improved, and the position needs to be calculated based on the slope correction.
            /*List<RotatedBox> initList = new ArrayList<>();
            for (RotatedBox result : detections) {
                // put low Y value at the head of the queue.
                initList.add(result);
            }
            Collections.sort(initList);

            List<ArrayList<RotatedBoxCompX>> lines = new ArrayList<>();
            List<RotatedBoxCompX> line = new ArrayList<>();
            RotatedBoxCompX firstBox = new RotatedBoxCompX(initList.get(0).getBox(), initList.get(0).getText());
            line.add(firstBox);
            lines.add((ArrayList) line);
            for (int i = 1; i < initList.size(); i++) {
                RotatedBoxCompX tmpBox = new RotatedBoxCompX(initList.get(i).getBox(), initList.get(i).getText());
                float y1 = firstBox.getBox().toFloatArray()[1];
                float y2 = tmpBox.getBox().toFloatArray()[1];
                float dis = Math.abs(y2 - y1);
                if (dis < 20) { // 认为是同 1 行  - Considered to be in the same line
                    line.add(tmpBox);
                } else { // 换行 - Line break
                    firstBox = tmpBox;
                    Collections.sort(line);
                    line = new ArrayList<>();
                    line.add(firstBox);
                    lines.add((ArrayList) line);
                }
            }

            for (int i = 0; i < lines.size(); i++) {
                ArrayList<RotatedBoxCompX> rotatedBoxCompXES = lines.get(i);
                for (int j = 0; j < rotatedBoxCompXES.size(); j++) {
                    RotatedBoxCompX rotatedBoxCompX = rotatedBoxCompXES.get(j);
                    // 边框
                    NDArray box = rotatedBoxCompX.getBox();
                    // 8个数值
                    float[] pointsArr = box.toFloatArray();

                    // 左上
                    float[] lt = java.util.Arrays.copyOfRange(pointsArr, 0, 2);

                    // 左下
                    float[] lb = java.util.Arrays.copyOfRange(pointsArr, 6, 8);

                    // 右上
                    float[] rt = java.util.Arrays.copyOfRange(pointsArr, 2, 4);

                    // 右下
                    float[] rb = java.util.Arrays.copyOfRange(pointsArr, 4, 6);
                    float width = distance(lt, rt);
                    float height = distance(lt, lb);

                    // 文字内容
                    String text = rotatedBoxCompX.getText();

                    PaddleWidget widget = new PaddleWidget();
                    widget.setText(text);
                    widget.setX1(lt[0]);
                    widget.setY1(lt[1]);
                    widget.setX2(rb[0]);
                    widget.setY2(rb[1]);
                    widget.setWidth(width);
                    widget.setHeight(height);
                    widgetList.add(widget);
                }
            }*/
        } catch (Exception e) {
            logger.error("ocr识图出错：" + e);
        }

        return widgetList;
    }

    /**
     * 释放 OCR paddle
     */
    private static void releasePaddle() {
        if (manager != null) {
            manager.close();
        }
        if (recognizer != null) {
            recognizer.close();
        }
        if (recognitionModel != null) {
            recognitionModel.close();
        }
        if (detector != null) {
            detector.close();
        }
        if (detectionModel != null) {
            detectionModel.close();
        }
    }

    private static float distance(float[] point1, float[] point2) {
        float disX = point1[0] - point2[0];
        float disY = point1[1] - point2[1];
        float dis = (float) Math.sqrt(disX * disX + disY * disY);
        return dis;
    }

    private static List<Widget> convertTrTextRes(String ocrContent) {
        List<Widget> trWidgetList = new ArrayList<>();
        JSONObject jsonObject = JSON.parseObject(ocrContent);
        Integer code = jsonObject.getInteger("code");
        if (code == HttpStatus.HTTP_OK) {
            JSONObject data = jsonObject.getJSONObject("data");
            JSONArray rawOutArray = data.getJSONArray("raw_out");
            for (int i = 0; i < rawOutArray.size(); i++) {
                TrWidget trWidget = new TrWidget();
                JSONArray jsonArrayRes = rawOutArray.getJSONArray(i);
                JSONArray pointArray = jsonArrayRes.getJSONArray(0);
                String text = jsonArrayRes.getString(1);
                BigDecimal confidence = jsonArrayRes.getBigDecimal(2);
                BigDecimal x = pointArray.getBigDecimal(0);
                BigDecimal y = pointArray.getBigDecimal(1);
                BigDecimal width = pointArray.getBigDecimal(2);
                BigDecimal height = pointArray.getBigDecimal(3);
                trWidget.setX(x);
                trWidget.setY(y);
                trWidget.setWidth(width);
                trWidget.setHeight(height);
                trWidget.setText(text);
                trWidget.setConfidence(confidence);
                trWidgetList.add(trWidget);
            }
        }
        return trWidgetList;
    }

    private static BufferedImage toBufferedImage(Image img) {
        if (img instanceof BufferedImage) {
            return (BufferedImage) img;
        }

        // Create a buffered image with transparency
        BufferedImage bimage = new BufferedImage(img.getWidth(null), img.getHeight(null), BufferedImage.TYPE_INT_ARGB);

        // Draw the image on to the buffered image
        Graphics2D bGr = bimage.createGraphics();
        bGr.drawImage(img, 0, 0, null);
        bGr.dispose();

        // Return the buffered image
        return bimage;
    }
}
